use anyhow::{anyhow, Result};
use clap::Parser;
use rmcp::{
    ServiceExt,
    model::{CallToolRequestParam, ClientCapabilities, ClientInfo, Implementation},
    transport::{StreamableHttpClientTransport, streamable_http_client::StreamableHttpClientTransportConfig},
};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

#[derive(Debug, Parser)]
#[command(
    name = "mcpdevtool",
    version,
    about = "MCP 命令行工具：连接 Streamable HTTP 服务器并调用工具或列出工具",
    long_about = r#"用途
—— 连接支持 Streamable HTTP 的 MCP 服务器，调用指定工具（输出 JSON CallToolResult），或列出服务器提供的工具。

服务器
—— 示例服务器（counter）位于 rust-sdk-main/examples/servers：
   1) 启动：cargo run --example servers_counter_streamhttp
   2) 端点：http://127.0.0.1:8000/mcp

用法示例
—— 列出工具：
   mcpdevtool --server-url http://127.0.0.1:8000/mcp --list-tools
—— 调用 increment（无参）：
   mcpdevtool --server-url http://127.0.0.1:8000/mcp --tool increment --args '{}'
—— 调用 say_hello（字符串参数）：
   mcpdevtool --server-url http://127.0.0.1:8000/mcp --tool say_hello --args '{"name":"Alice"}'
—— 调用 sum（数字参数）：
   mcpdevtool --server-url http://127.0.0.1:8000/mcp --tool sum --args '{"a":3,"b":4}'

参数说明
—— --args 必须是 JSON 对象；无参场景请传 '{}'
—— --args 与 --tool-args 等价（别名）。
—— --token 可选，若服务器需要认证，将作为 Bearer 令牌加入 Authorization 头。

日志
—— 通过 RUST_LOG 控制日志级别，例如：RUST_LOG=info mcpdevtool ...
"#
)]
struct Args {
    /// MCP 服务器 URL，例如：http://localhost:8000/mcp
    #[arg(long, short = 'u', value_name = "URL")]
    server_url: String,

    /// 列出服务器提供的工具
    #[arg(long)]
    list_tools: bool,

    /// 仅输出工具名列表（与 --tool 互斥）
    #[arg(long)]
    list_names: bool,

    /// 要调用的 MCP 工具名（与 --list-tools/--list-names 互斥）
    #[arg(long, short = 't', value_name = "TOOL", required_unless_present_any = ["list_tools", "list_names"])]
    tool: Option<String>,

    /// 工具参数（JSON 字符串，必须是对象），例如：'{"a":1,"b":2}'
    #[arg(long = "args", short = 'a', alias = "tool-args", value_name = "JSON")]
    tool_args: Option<String>,

    /// 可选的 Bearer Token（作为 Authorization 头）
    #[arg(long)]
    token: Option<String>,
}

#[tokio::main]
async fn main() -> Result<()> {
    // 初始化日志
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| format!("info,{}=debug", env!("CARGO_CRATE_NAME")).into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    let args = Args::parse();

    // 创建基于 Streamable HTTP 的客户端传输（支持 Authorization 头）
    let config = if let Some(tok) = &args.token {
        StreamableHttpClientTransportConfig::with_uri(args.server_url.clone())
            .auth_header(format!("Bearer {}", tok))
    } else {
        StreamableHttpClientTransportConfig::with_uri(args.server_url.clone())
    };
    let transport = StreamableHttpClientTransport::from_config(config);

    // 客户端信息（可根据需要调整 capabilities）
    let client_info = ClientInfo {
        protocol_version: Default::default(),
        capabilities: ClientCapabilities::default(),
        client_info: Implementation {
            name: "mcpdevtool".to_string(),
            version: env!("CARGO_PKG_VERSION").to_string(),
        },
    };

    // 建立连接
    let client = client_info
        .serve(transport)
        .await
        .inspect_err(|e| tracing::error!("client error: {:?}", e))?;

    // 输出服务器信息
    let server_info = client.peer_info();
    tracing::info!("Connected to server: {server_info:#?}");

    if args.list_tools || args.list_names {
        // 列出工具
        let tools = client.list_tools(Default::default()).await?;
        if args.list_names {
            // 仅输出工具名数组
            let names: Vec<String> = tools.tools.into_iter().map(|t| t.name.to_string()).collect();
            println!("{}", serde_json::to_string_pretty(&names)?);
        } else {
            println!("{}", serde_json::to_string_pretty(&tools)?);
        }
        client.cancel().await?;
        return Ok(());
    }

    // 解析 tool_args 为 JSON 对象（仅在调用工具模式下）
    let arguments = match args.tool_args {
        Some(s) => {
            let v: serde_json::Value = serde_json::from_str(&s)?;
            match v {
                serde_json::Value::Object(map) => Some(map),
                serde_json::Value::Null => None,
                _ => return Err(anyhow!(r#"tool-args 必须是 JSON 对象 (例如: {{"key":"value"}})"#)),
            }
        }
        None => None,
    };

    // 调用指定工具
    let tool_name = args.tool.expect("tool 参数缺失：当未使用 --list-tools 时必须提供 --tool");
    let result = client
        .call_tool(CallToolRequestParam {
            name: tool_name.into(),
            arguments,
        })
        .await?;

    // 打印结果为 JSON
    println!("{}", serde_json::to_string_pretty(&result)?);

    // 关闭连接
    client.cancel().await?;
    Ok(())
}